package com.wang.springbootinit.sparkAi;

import cn.hutool.core.util.StrUtil;
import com.wang.springbootinit.common.ErrorCode;
import com.wang.springbootinit.exception.BusinessException;
import io.github.briqt.spark4j.SparkClient;
import io.github.briqt.spark4j.constant.SparkApiVersion;
import io.github.briqt.spark4j.exception.SparkException;
import io.github.briqt.spark4j.model.SparkMessage;
import io.github.briqt.spark4j.model.SparkSyncChatResponse;
import io.github.briqt.spark4j.model.request.SparkRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * 抽象Ai分析处理器
 */

@Component
@Slf4j
public class AiAnalyseHandler {
    @Resource
    private SparkClient sparkClient;

    /**
     * AI模型id
     */
    private final static Long modelId = 1675876907536691202L;

    /**
     * 同步调用
     *
     * @param analyseQueryParam
     * @return
     */
    public AiAnalyseModel syncChat(AiAnalyseQueryParam analyseQueryParam) {
        try {
            SparkRequest sparkRequest = this.getSparkRequest(analyseQueryParam);

            SparkSyncChatResponse chatResponse = sparkClient.chatSync(sparkRequest);
            if (chatResponse.isOk()) {
                return new AiAnalyseModel(chatResponse);
            }
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, chatResponse.getException().toString());
        } catch (SparkException e) {
            log.error("AiAnalyseHandler syncChat exception ", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, ErrorCode.SYSTEM_ERROR.getMessage());
        }
    }

    /**
     * 流式调用
     *
     * @param analyseQueryParam
     */
    public void chatStream(AiAnalyseQueryParam analyseQueryParam) {
        try {
            SparkRequest sparkRequest = this.getSparkRequest(analyseQueryParam);
            sparkClient.chatStream(sparkRequest, new SparkConsoleListener());
        } catch (Exception e) {
            log.error("AiAnalyseHandler chatStream exception ", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, ErrorCode.SYSTEM_ERROR.getMessage());
        }
    }


    /**
     * 建单调用
     * @param analyseQueryParam
     * @return
     */
    public AiAnalyseModel simpleChat(AiAnalyseQueryParam analyseQueryParam) {
        try {
            SparkRequest sparkRequest = this.getSparkRequest(analyseQueryParam,analyseQueryParam.getSystemContent());
            SparkSyncChatResponse chatResponse = sparkClient.chatSync(sparkRequest);
            if (chatResponse.isOk()) {
                return new AiAnalyseModel(chatResponse);
            }
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, chatResponse.getException().toString());
        } catch (SparkException e) {
            log.error("AiAnalyseHandler syncChat exception ", e);
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, ErrorCode.SYSTEM_ERROR.getMessage());
        }
    }

    private SparkRequest getSparkRequest(AiAnalyseQueryParam analyseQueryParam) {
        List<SparkMessage> messages = new ArrayList<>();
        String systemContent = StrUtil.isEmpty(analyseQueryParam.getSystemContent()) ? AiAnalyseConstant.SystemContent : analyseQueryParam.getSystemContent();
        messages.add(SparkMessage.systemContent(systemContent));
        messages.add(SparkMessage.userContent(analyseQueryParam.getUserQuestionContent()));

        return SparkRequest.builder()
                .messages(messages)
                .maxTokens(AiAnalyseConstant.maxTokens)
                .temperature(AiAnalyseConstant.temperature)
                .apiVersion(SparkApiVersion.V3_0)
                .build();

    }

    private SparkRequest getSparkRequest(AiAnalyseQueryParam analyseQueryParam,String systemContent) {
        List<SparkMessage> messages = new ArrayList<>();
        if(StrUtil.isNotEmpty(systemContent)){
            messages.add(SparkMessage.systemContent(systemContent));
        }
        messages.add(SparkMessage.userContent(analyseQueryParam.getUserQuestionContent()));

        return SparkRequest.builder()
                .messages(messages)
                .maxTokens(AiAnalyseConstant.maxTokens)
                .temperature(AiAnalyseConstant.temperature)
                .apiVersion(SparkApiVersion.V3_0)
                .build();

    }
}
